Use TPUs  

This guide demonstrates how to perform basic training on Tensor Processing Units (TPUs) and TPU Pods, a collection of TPU devices connected by dedicated high-speed network interfaces, with tf.keras and custom training loops.

TPUs are Google's custom-developed application-specific integrated circuits (ASICs) used to accelerate machine learning workloads. They are available through Google Colab, the TPU Research Cloud, and Cloud TPU.

Before you run this Colab notebook, make sure that your hardware accelerator is a TPU by checking your notebook settings: Runtime > Change runtime type > Hardware accelerator > TPU.

Import some necessary libraries, including TensorFlow Datasets:

TPUs are typically Cloud TPU workers, which are different from the local process running the user's Python program. Thus, you need to do some initialization work to connect to the remote cluster and initialize the TPUs. Note that the tpu argument to tf.distribute.cluster_resolver.TPUClusterResolver is a special address just for Colab. If you are running your code on Google Compute Engine (GCE), you should instead pass in the name of your Cloud TPU.

After the TPU is initialized, you can use manual device placement to place the computation on a single TPU device:

Usually, you run your model on multiple TPUs in a data-parallel way. To distribute your model on multiple TPUs (as well as multiple GPUs or multiple machines), TensorFlow offers the tf.distribute.Strategy API. You can replace your distribution strategy and the model will run on any given (TPU) device. Learn more in the Distributed training with TensorFlow guide.

Using the tf.distribute.TPUStrategy option implements synchronous distributed training. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores, which are used in TPUStrategy.

To demonstrate this, create a tf.distribute.TPUStrategy object:

To replicate a computation so it can run in all TPU cores, you can pass it into the Strategy.run API. Below is an example that shows all cores receiving the same inputs (a, b) and performing matrix multiplication on each core independently. The outputs will be the values from all the replicas.

Having covered the basic concepts, consider a more concrete example. This section demonstrates how to use the distribution strategy—tf.distribute.TPUStrategy—to train a Keras model on a Cloud TPU.

Start with a definition of a Sequential Keras model for image classification on the MNIST dataset. It's no different than what you would use if you were training on CPUs or GPUs. Note that Keras model creation needs to be inside the Strategy.scope, so the variables can be created on each TPU device. Other parts of the code are not necessary to be inside the Strategy scope.

Efficient use of the tf.data.Dataset API is critical when using a Cloud TPU. You can learn more about dataset performance in the Input pipeline performance guide.

If you are using TPU Nodes, you need to store all data files read by the TensorFlow Dataset in Google Cloud Storage (GCS) buckets. If you are using TPU VMs, you can store data wherever you like. For more information on TPU Nodes and TPU VMs, refer to the TPU System Architecture documentation.

For most use cases, it is recommended to convert your data into the TFRecord format and use a tf.data.TFRecordDataset to read it. Check the TFRecord and tf.Example tutorial for details on how to do this. It is not a hard requirement and you can use other dataset readers, such as tf.data.FixedLengthRecordDataset or tf.data.TextLineDataset.

You can load entire small datasets into memory using tf.data.Dataset.cache.

Regardless of the data format used, it is strongly recommended that you use large files on the order of 100MB. This is especially important in this networked setting, as the overhead of opening a file is significantly higher.

As shown in the code below, you should use the Tensorflow Datasets tfds.load module to get a copy of the MNIST training and test data. Note that try_gcs is specified to use a copy that is available in a public GCS bucket. If you don't specify this, the TPU will not be able to access the downloaded data.

You can train your model with Keras Model.fit and Model.compile APIs. There is nothing TPU-specific in this step—you write the code as if you were using multiple GPUs and a MirroredStrategy instead of the TPUStrategy. You can learn more in the Distributed training with Keras tutorial.

To reduce Python overhead and maximize the performance of your TPU, pass in the steps_per_execution argument to Keras Model.compile. In this example, it increases throughput by about 50%:

You can also create and train your model using tf.function and tf.distribute APIs directly. You can use the Strategy.experimental_distribute_datasets_from_function API to distribute the tf.data.Dataset given a dataset function. Note that in the example below the batch size passed into the Dataset is the per-replica batch size instead of the global batch size. To learn more, check out the Custom training with tf.distribute.Strategy tutorial.

First, create the model, datasets and tf.functions:

Then, run the training loop:

You can improve the performance by running multiple steps within a tf.function. This is achieved by wrapping the Strategy.run call with a tf.range inside tf.function, and AutoGraph will convert it to a tf.while_loop on the TPU worker. You can learn more about tf.functions in the Better performance with tf.function guide.

Despite the improved performance, there are tradeoffs with this method compared to running a single step inside a tf.function. Running multiple steps in a tf.function is less flexible—you cannot run things eagerly or arbitrary Python code within the steps.

To learn more about Cloud TPUs and how to use them: